"""
labelme2yolo_detect.py
create by cuiyu 2025-08-04
说明：
 * 将 labelme 的 对象检测 任务标注（矩形框）转换为 yolo 格式

用法:
  labelme2yolo_detectg.py

更新:
2025-08-04:
"""

import json
import os
from pathlib import Path

class _Labelme2Yolo_Detect:

    def l2y_file(self, json_file: Path, label_map: dict):
        """
        将单个 LabelMe JSON 文件 (COCO 格式) 转换为YOLOv11格式的TXT标注文件，保存在同一个目录下
        
        :param json_file: LabelMe JSON文件路径
        :param label_map: 类别名称到ID的映射字典（如{"car":0, "bus":1}）
        """
        json_file = Path(json_file)
        
        # 读取JSON文件
        with open(json_file, 'r', encoding='utf-8') as f:
            labelme_data = json.load(f)
        
        # 获取图像尺寸（必须存在）
        image_width = labelme_data['imageWidth']
        image_height = labelme_data['imageHeight']
        
        # 存储YOLO格式的每行标注
        yolo_lines = []
        
        # 遍历所有标注形状
        for shape in labelme_data['shapes']:
            # 只处理矩形框标注
            if shape['shape_type'] != 'rectangle':
                continue
            
            label = shape['label']
            # 跳过未定义类别的标签
            if label not in label_map:
                print(f"⚠️ 警告: 跳过未定义标签 '{label}'，请添加到label_map")
                continue
            
            class_id = label_map[label]
            points = shape['points']
            
            # 确保矩形有两个点（左上和右下）
            if len(points) != 2:
                print(f"⚠️ 警告: {json_file} 中矩形标注点数错误")
                continue
            
            # 提取坐标并确保顺序（左上->右下）
            x1, y1 = points[0]
            x2, y2 = points[1]
            
            # 转换为实际坐标范围（处理反向标注）
            x_min = min(x1, x2)
            x_max = max(x1, x2)
            y_min = min(y1, y2)
            y_max = max(y1, y2)
            
            # 计算归一化中心坐标和宽高（YOLO格式核心）
            x_center = (x_min + x_max) / 2.0 / image_width
            y_center = (y_min + y_max) / 2.0 / image_height
            width = (x_max - x_min) / image_width
            height = (y_max - y_min) / image_height
            
            # 确保坐标在[0,1]范围内（防止标注越界）
            x_center = max(0, min(1, x_center))
            y_center = max(0, min(1, y_center))
            width = max(0, min(1, width))
            height = max(0, min(1, height))
            
            # 格式: <class_id> <x_center> <y_center> <width> <height>
            yolo_line = f"{class_id} {x_center:.6f} {y_center:.6f} {width:.6f} {height:.6f}"
            yolo_lines.append(yolo_line)
        
        # 写入输出文件（与JSON同名的TXT）
        txt_filename = json_file.with_suffix(".txt")
        
        with open(txt_filename, 'w') as f:
            f.write("\n".join(yolo_lines))

    def l2y_dir(self, input_dir: Path, label_map: dict):
        """
        批量转换整个目录下的LabelMe JSON文件，生成的 txt 保存在同一个目录，与输入JSON文件同名的TXT
        
        :param input_dir: 包含JSON文件的目录
        :param label_map: 类别映射字典, 如 {"label": index} 格式
        """

        input_dir = Path(input_dir)

        json_files = [Path(file) for file in input_dir.rglob("*.json")]        
        num = len(json_files)
        n_len = len(str(num))
        for i, json_file in enumerate(json_files):
            n = i+1
            print(f'{n:{n_len}} / {num} {float(n)/num*100:.2f}%  file={json_file.name}', end='\r')
            self.l2y_file(json_file, label_map)
        print("")


        print(f"✅ 转换完成！共处理 {len(json_files)} 个文件。")

Coco2Yolo = _Labelme2Yolo_Detect()

def _main():
    from bhclasses_define import BHClasses
    # 配置参数
    # LABEL_MAP = {
    #     "car": 0,    # 确保与训练配置一致
    #     "bus": 1,
    #     "person": 2
    # }
    
    INPUT_JSON_DIR = Path("/aaby/netdisk/mydata/ai-data/dataset/priv/井盖病害/训练数据/img01")  # LabelMe JSON文件目录
    
    # 执行转换
    Coco2Yolo.l2y_dir(INPUT_JSON_DIR, BHClasses.dict_name_id)

# ====== 使用示例 ======
if __name__ == "__main__":
    _main()

